#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
从图片中提取LoRA信息
"""

import os
import sys
import json
from PIL import Image
from PIL.PngImagePlugin import PngInfo

# 添加当前目录到路径
current_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
sys.path.insert(0, current_dir)

def extract_lora_from_workflow(workflow_data):
    """从工作流数据中提取LoRA信息"""
    lora_info = []
    
    try:
        if isinstance(workflow_data, str):
            workflow_json = json.loads(workflow_data)
        else:
            workflow_json = workflow_data
            
        print(f"  工作流节点总数: {len(workflow_json.get('nodes', []))}")
        
        nodes = workflow_json.get('nodes', {})
        
        # 处理不同的节点格式
        if isinstance(nodes, dict):
            node_items = nodes.items()
        elif isinstance(nodes, list):
            node_items = enumerate(nodes)
        else:
            print(f"  [ERROR] 未知的节点格式: {type(nodes)}")
            return lora_info
        
        for node_id, node_data in node_items:
            class_type = node_data.get('class_type', '')
            inputs = node_data.get('inputs', {})
            
            # 检查是否是LoRA相关节点（扩展关键词）
            lora_keywords = ['lora', 'LoRA', 'LORA', 'lora_loader', 'LoraLoader', 'lora_apply', 'LoraApply']
            is_lora_node = any(keyword in class_type for keyword in lora_keywords)
            
            # 也检查输入参数中是否有LoRA相关字段
            lora_inputs = ['lora_name', 'lora', 'model', 'lora_model', 'lora_path']
            has_lora_inputs = any(key in inputs for key in lora_inputs)
            
            # 检查PowerLoraLoaderHeaderWidget节点
            has_power_lora = 'PowerLoraLoaderHeaderWidget' in inputs
            
            # 检查widgets_values中的PowerLoraLoaderHeaderWidget
            widgets_values = node_data.get('widgets_values', [])
            has_power_lora_widget = False
            if isinstance(widgets_values, list):
                for widget in widgets_values:
                    if isinstance(widget, dict) and widget.get('type') == 'PowerLoraLoaderHeaderWidget':
                        has_power_lora_widget = True
                        break
            
            if is_lora_node or has_lora_inputs or has_power_lora or has_power_lora_widget:
                lora_node = {
                    'node_id': str(node_id),
                    'class_type': class_type,
                    'inputs': inputs,
                    'is_power_lora': has_power_lora
                }
                
                if has_power_lora or has_power_lora_widget:
                    # 处理PowerLoraLoaderHeaderWidget节点
                    print(f"    找到PowerLoraLoaderHeaderWidget节点: {class_type}")
                    
                    # 提取所有激活的LoRA
                    active_loras = []
                    
                    # 从inputs中提取
                    if isinstance(inputs, dict):
                        for key, value in inputs.items():
                            if key.startswith('lora_') and isinstance(value, dict):
                                lora_name = value.get('lora', '')
                                is_on = value.get('on', False)
                                strength = value.get('strength', 1.0)
                                
                                if lora_name and is_on:
                                    active_loras.append({
                                        'name': lora_name,
                                        'strength': strength
                                    })
                                    print(f"      激活的LoRA: {lora_name} (强度: {strength})")
                    
                    # 从widgets_values中提取
                    if isinstance(widgets_values, list):
                        for widget in widgets_values:
                            if isinstance(widget, dict) and 'lora' in widget:
                                lora_name = widget.get('lora', '')
                                is_on = widget.get('on', False)
                                strength = widget.get('strength', 1.0)
                                
                                if lora_name and is_on:
                                    # 避免重复
                                    if not any(lora['name'] == lora_name for lora in active_loras):
                                        active_loras.append({
                                            'name': lora_name,
                                            'strength': strength
                                        })
                                        print(f"      激活的LoRA: {lora_name} (强度: {strength})")
                    
                    lora_node['active_loras'] = active_loras
                    lora_node['lora_count'] = len(active_loras)
                    
                else:
                    # 处理传统LoRA节点
                    # 提取LoRA路径
                    lora_path = None
                    if 'lora_name' in inputs:
                        lora_path = inputs['lora_name']
                    elif 'model' in inputs:
                        lora_path = inputs['model']
                    elif 'lora' in inputs:
                        lora_path = inputs['lora']
                    
                    if lora_path:
                        lora_node['lora_path'] = lora_path
                    
                    # 提取强度信息
                    strength = None
                    if 'strength_model' in inputs:
                        strength = inputs['strength_model']
                    elif 'strength' in inputs:
                        strength = inputs['strength']
                    
                    if strength is not None:
                        lora_node['strength'] = strength
                    
                    print(f"    找到LoRA节点: {class_type}")
                    print(f"      LoRA路径: {lora_path}")
                    print(f"      强度: {strength}")
                
                lora_info.append(lora_node)
        
        return lora_info
        
    except Exception as e:
        print(f"  [ERROR] 解析工作流失败: {e}")
        return lora_info

def extract_lora_from_image(image_path):
    """从图片中提取LoRA信息"""
    try:
        print(f"[INFO] 分析图片: {image_path}")
        
        # 打开图片
        with Image.open(image_path) as img:
            print(f"  图片尺寸: {img.size}")
            print(f"  图片模式: {img.mode}")
            
            # 检查PNG信息
            if hasattr(img, 'info') and img.info:
                print(f"  PNG信息键: {list(img.info.keys())}")
                
                # 检查工作流信息
                if 'workflow' in img.info:
                    workflow_data = img.info['workflow']
                    print(f"  工作流数据类型: {type(workflow_data)}")
                    
                    if isinstance(workflow_data, str):
                        print(f"  工作流数据长度: {len(workflow_data)} 字符")
                        
                        # 提取LoRA信息
                        lora_info = extract_lora_from_workflow(workflow_data)
                        
                        if lora_info:
                            print(f"\n[SUCCESS] 找到 {len(lora_info)} 个LoRA节点:")
                            for i, lora in enumerate(lora_info, 1):
                                print(f"  LoRA {i}:")
                                print(f"    节点ID: {lora['node_id']}")
                                print(f"    节点类型: {lora['class_type']}")
                                print(f"    LoRA路径: {lora.get('lora_path', '未知')}")
                                print(f"    强度: {lora.get('strength', '未知')}")
                                print(f"    完整输入: {lora['inputs']}")
                                print()
                            
                            return lora_info
                        else:
                            print("  [INFO] 未找到LoRA节点")
                    else:
                        print(f"  [ERROR] 工作流数据类型错误: {type(workflow_data)}")
                else:
                    print("  [ERROR] 图片中没有工作流信息")
                    
                # 检查其他元数据
                print(f"\n  其他元数据:")
                for key, value in img.info.items():
                    if key not in ['workflow']:
                        value_str = str(value)
                        if len(value_str) > 200:
                            value_str = value_str[:200] + "..."
                        print(f"    {key}: {value_str}")
                
                # 输出完整的工作流数据
                print(f"\n  完整工作流数据:")
                print(f"    {workflow_data[:500]}...")
                
                # 尝试解析工作流并显示所有节点类型
                try:
                    workflow_json = json.loads(workflow_data)
                    nodes = workflow_json.get('nodes', {})
                    
                    print(f"\n  所有节点类型:")
                    node_types = set()
                    if isinstance(nodes, dict):
                        for node_id, node_data in nodes.items():
                            class_type = node_data.get('class_type', '')
                            node_types.add(class_type)
                    elif isinstance(nodes, list):
                        for node_data in nodes:
                            class_type = node_data.get('class_type', '')
                            node_types.add(class_type)
                    
                    for node_type in sorted(node_types):
                        print(f"    - {node_type}")
                    
                    # 显示所有节点的详细信息
                    print(f"\n  所有节点详细信息:")
                    if isinstance(nodes, dict):
                        for node_id, node_data in nodes.items():
                            class_type = node_data.get('class_type', '')
                            inputs = node_data.get('inputs', {})
                            print(f"    节点 {node_id}: {class_type}")
                            for key, value in inputs.items():
                                if isinstance(value, str) and len(value) > 50:
                                    value = value[:50] + "..."
                                print(f"      {key}: {value}")
                    elif isinstance(nodes, list):
                        for i, node_data in enumerate(nodes):
                            class_type = node_data.get('class_type', '')
                            inputs = node_data.get('inputs', {})
                            print(f"    节点 {i}: {class_type}")
                            for key, value in inputs.items():
                                if isinstance(value, str) and len(value) > 50:
                                    value = value[:50] + "..."
                                print(f"      {key}: {value}")
                    
                    # 保存完整工作流到文件
                    workflow_file = os.path.join(current_dir, "scripts", "full_workflow.json")
                    try:
                        with open(workflow_file, 'w', encoding='utf-8') as f:
                            json.dump(workflow_json, f, ensure_ascii=False, indent=2)
                        print(f"\n  [INFO] 完整工作流已保存到: {workflow_file}")
                    except Exception as e:
                        print(f"\n  [ERROR] 保存工作流失败: {e}")
                        
                except Exception as e:
                    print(f"    [ERROR] 解析工作流失败: {e}")
            else:
                print("  [ERROR] 图片没有元数据信息")
                
        return []
        
    except Exception as e:
        print(f"[ERROR] 分析图片失败: {e}")
        import traceback
        traceback.print_exc()
        return []

def save_lora_info_to_file(lora_info, output_file):
    """保存LoRA信息到文件"""
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(lora_info, f, ensure_ascii=False, indent=2)
        print(f"[SUCCESS] LoRA信息已保存到: {output_file}")
    except Exception as e:
        print(f"[ERROR] 保存文件失败: {e}")

def main():
    """主函数"""
    print("=" * 60)
    print("LoRA信息提取工具")
    print("=" * 60)
    
    # 测试图片路径
    image_path = os.path.join(current_dir, "1.png")
    output_file = os.path.join(current_dir, "scripts", "extracted_lora_info.json")
    
    if not os.path.exists(image_path):
        print(f"[ERROR] 图片文件不存在: {image_path}")
        return
    
    print(f"[INFO] 开始分析图片: {image_path}")
    print()
    
    # 提取LoRA信息
    lora_info = extract_lora_from_image(image_path)
    
    print("=" * 60)
    print("提取结果总结")
    print("=" * 60)
    
    if lora_info:
        print(f"[SUCCESS] 成功提取到 {len(lora_info)} 个LoRA节点")
        
        # 显示所有LoRA路径
        print("\n所有LoRA路径:")
        for i, lora in enumerate(lora_info, 1):
            lora_path = lora.get('lora_path', '未知')
            strength = lora.get('strength', '未知')
            print(f"  {i}. {lora_path} (强度: {strength})")
        
        # 保存到文件
        save_lora_info_to_file(lora_info, output_file)
        
        # 显示详细信息
        print(f"\n详细信息:")
        for i, lora in enumerate(lora_info, 1):
            print(f"\nLoRA {i}:")
            print(f"  节点ID: {lora['node_id']}")
            print(f"  节点类型: {lora['class_type']}")
            print(f"  LoRA路径: {lora.get('lora_path', '未知')}")
            print(f"  强度: {lora.get('strength', '未知')}")
            print(f"  完整输入参数:")
            inputs = lora.get('inputs', {})
            if isinstance(inputs, dict):
                for key, value in inputs.items():
                    print(f"    {key}: {value}")
            else:
                print(f"    输入类型: {type(inputs)}")
                print(f"    输入内容: {inputs}")
    else:
        print("[INFO] 未找到LoRA信息")
    
    print("\n" + "=" * 60)
    print("分析完成")

if __name__ == "__main__":
    main()
